import json
import random
import torch
import numpy as np
import pdb
import jieba
import jieba.posseg as psg
import math


bio2id = {'O': 0, 'B': 1, 'I': 2, 'E': 3, 'S': 4}



class DataLoader(object):
    """
    Load data from json files, preprocess and prepare batches.
    """
    def __init__(self, args, file_path):
        self.batch_size = args.batch_size
        self.args = args
        self.file_path = file_path
        self.word2id = {w:i for i,w in enumerate(args.word_vocab)}
        self.raw_data = []

        with open(file_path, 'r', encoding='utf8') as f:
            #self.raw_data = json.load(f)
            for line in f.readlines():
                if line.strip() != '':
                    self.raw_data.append(json.loads(line))
        self.pos2id, self.label2id = args.pos2id, args.label2id
        self.char2id = args.char2id
        self.data = self.preprocess(self.raw_data)
        self.num_examples = len(self.data)
        #pdb.set_trace()
        # chunk into batches
        self.data = [self.data[i:i+self.batch_size] for i in range(0, len(self.data), self.batch_size)]
        print("{} batches created for {}".format(len(self.data), self.file_path))

    def gen_labels(self, sentText, relationMentions):
        #sent_words = nltk.word_tokenize(sentText)
        sent_words = list(sentText)
        l = len(sent_words)

        labels = np.zeros((l,l,len(self.args.label2id)), dtype=np.float32)
        for i, rm in enumerate(relationMentions):
            #pdb.set_trace()
            em1 = rm['em1Text']
            em2 = rm['em2Text']

            if find_index(sent_words, em1) is not None and find_index(sent_words, em2) is not None:
                e1_sta, e1_end = find_index(sent_words, em1)
                e2_sta, e2_end = find_index(sent_words, em2)

                #h_e1, h_e2 = sent_words.index(rm['em1Text'][0]), sent_words.index(rm['em2Text'][0])
                #labels[h_e1][h_e2][self.args.label2id[rm['label']]] = 1
                labels[e1_sta][e2_sta][self.args.label2id[rm['label']]] = 1
                #pdb.set_trace()
            else:
                continue
        return labels

    def find_pos(self, sent_list, word_list):
        '''
        return position list
        '''
        l = len(word_list)
        for i in range(len(sent_list)):
            flag = True
            j = 0
            while j < l:
                if word_list[j] != sent_list[i + j]:
                    flag = False
                    break
                j += 1
            if flag:
                return range(i, i+l)
        return []

    def padding_labels(self, labels, batch_size):
        """ Convert labels to a padded LongTensor. """
        token_len = max(x.shape[0] for x in labels)
        padded_labels = torch.FloatTensor(batch_size, token_len, token_len, len(self.label2id)).fill_(0)
        for i, s in enumerate(labels):
            padded_labels[i,:s.shape[0],:s.shape[0]] = torch.FloatTensor(s)
        return padded_labels

    def preprocess(self, data):
        """ Preprocess the data and convert to ids. """
        processed = []

        for d in data:

            sent = d['sentText'].lower()
            #pos    = d['pos']

            breakflag = False

            seglist = psg.cut(sent)
            sent_words = []     #因为中文命名实体识别，因此直接按照字符级去捕获
            for char in sent:
                sent_words.append(char)

            # sent_pos = []
            #
            # for seg in seglist:
            #
            #     sent_words.append(seg.word)
            #
            #     sent_pos.append(seg.flag)

            # sent_pos = [ for w in sent_words]
            #
            # sent_words = nltk.word_tokenize(sent)
            # sent_pos = nltk.pos_tag(sent_words)


            # map to ids
            sent_chars = []

            radius = math.floor(self.args.kernel_radius/2)


            for i, char in enumerate(sent_words):       #每个字符各取其前后五个词
                chartoks = []
                for count in range(i-radius, i+radius):
                    #pdb.set_trace()
                    #print(count)
                    if count >=0 and count < len(sent_words):
                        tmpid = self.char2id.get(sent_words[count], 1)
                        chartoks.append(tmpid)
                    else:
                        chartoks.append(0)
                #pdb.set_trace()
                sent_chars.append(chartoks)
                #pdb.set_trace()


            # chartoks = [self.char2id.get(ctok, 1) for ctok in list(sent_words)]
            # word_len = min(len(chartoks), self.args.max_sent_len)
            # for _ in range(self.args.max_sent_len - word_len):
            #     chartoks.append(0)
            # sent_chars.append(chartoks[:self.args.max_sent_len])

            #pdb.set_trace()
            tokens = map_to_ids(sent_words, self.word2id)       #UNK过多会非常影响训练效果,英文大小写非常重要！
            pos = None
            #sent_chars = None

            ##!!!因为是修饰词，所以基于char去重新设计pos信息，名词与形
            # try:
            #     pos    = map_to_ids(sent_pos, self.pos2id)
            # except:
            #     print(sent_pos)
            #     print("pos appear UNK ")
            #     pos    = map_to_pos_ids(sent_pos, self.pos2id)

            l      = len(tokens)
            #pdb.set_trace()



            NER_labels = [bio2id['O'] for _ in range(l)]
            d['en_list'] = []
            for ent in d['entityMentions']:
                d['en_list'].append(ent['text'].lower())           #这里有一个问题：词向量中词典都是小写，但是实际实体中均为大写（大写可能蕴含有信息），但如果都用大写的话UNK过多会导致效果变差，因此改用小写

            for en in d['en_list']:
                #find_posi = self.find_pos(sent_words, en)

                if len(en) == 1:
                    en_index = sent_words.index(en[0])

                    NER_labels[en_index] = bio2id['S']
                else:
                    #此处可能会出现由分词导致的不对齐问题，可以直接跳过此种情况；['Bloomberg', 'L.P', '.']  ['Bloomberg', 'L.P.']
                    if find_index(sent_words, en) is not None:
                        sta, end = find_index(sent_words, en)
                        tmp_block = [bio2id['B']] + [bio2id['I'] for _ in range(end-sta-1)] + [bio2id['E']]
                        NER_labels[sta:end+1] = tmp_block
                    else:
                        breakflag = True
                        break
            if breakflag:
                continue


            # sentence mask
            mask_s = [1 for i in range(l)]

            #pdb.set_trace()
            RC_labels = self.gen_labels(d['sentText'], d['relationMentions'])
            processed += [(tokens,  pos, sent_chars, mask_s, (d['sentText'], d['relationMentions']), NER_labels)]

        #pdb.set_trace()
        return processed

    def __len__(self):
        return len(self.data)

    # 0: tokens, 1: mask_s, 2: label
    def __getitem__(self, key):
    #def getitem(self, key):
        """ Get a batch with index. """
        if not isinstance(key, int):
            raise TypeError
        if key < 0 or key >= len(self.data):
            raise IndexError
        batch = self.data[key]

        batch_size = len(batch)
        batch = list(zip(*batch))
        assert len(batch) == 6

        # sort all fields by lens for easy RNN operations
        lens = [len(x) for x in batch[0]]
        batch, _ = sort_all(batch, lens)

        # convert to tensors
        words = get_long_tensor(batch[0], batch_size)
        #pos   = get_long_tensor(batch[1], batch_size)

        pos = batch[1]

        char = get_long_tensor_char(batch[2], batch_size, self.args.max_word_len)



        #char = batch[2]
        
        # mask_s to tensors 
        mask_s = get_float_tensor(batch[3], batch_size)


        # RC_labels to tensors
        RC_labels = [self.gen_labels(sentText, relationMentions) for sentText, relationMentions in batch[4]]
        RC_labels = self.padding_labels(RC_labels,  batch_size)

        #pdb.set_trace()

        # NER_labels to tensors
        NER_labels = get_long_tensor(batch[5], batch_size)
        
        return [words, pos, char, mask_s, NER_labels, RC_labels]

    def __iter__(self):
        for i in range(self.__len__()):
            yield self.__getitem__(i)
            #yield self.getitem(i)

def find_index(sentText, en): # [sta, end]

    for i in range(len(sentText)):
        if sentText[i:len(en)+i] == list(en):
            return i, len(en)+i-1

def map_to_ids(tokens, vocab):
    ids = [vocab[t] if t in vocab else vocab['<UNK>'] for t in tokens]
    return ids

def map_to_pos_ids(tokens, vocab):
    ids = [vocab[t] if t in vocab else vocab['.'] for t in tokens]
    return ids

def get_long_tensor(tokens_list, batch_size):
    """ Convert list of list of tokens to a padded LongTensor. """
    token_len = max(len(x) for x in tokens_list)
    tokens = torch.LongTensor(batch_size, token_len).fill_(0)
    #pdb.set_trace()
    for i, s in enumerate(tokens_list):
        tokens[i, :len(s)] = torch.LongTensor(s)
    return tokens

def get_long_tensor_char(tokens_list, batch_size, max_word_len):
    """ Convert list of list of tokens to a padded LongTensor. """
    token_len = max(len(x) for x in tokens_list)
    tokens = torch.LongTensor(batch_size, token_len, max_word_len).fill_(0)
    #pdb.set_trace()
    for i, s in enumerate(tokens_list):
        tokens[i, :len(s), :] = torch.LongTensor(s)
    return tokens

def get_float_tensor(tokens_list, batch_size):
    """ Convert list of list of tokens to a padded FloatTensor. """
    token_len = max(len(x) for x in tokens_list)
    tokens = torch.FloatTensor(batch_size, token_len).fill_(0)
    for i, s in enumerate(tokens_list):
        tokens[i, :len(s)] = torch.FloatTensor(s)
    return tokens

def sort_all(batch, lens):
    """ Sort all fields by descending order of lens, and return the original indices. """
    unsorted_all = [lens] + [range(len(lens))] + list(batch)
    sorted_all = [list(t) for t in zip(*sorted(zip(*unsorted_all), reverse=True))]
    return sorted_all[2:], sorted_all[1]

